import argparse
import os
import pysam
import pod5
from pod5 import Writer as Pod5Writer
from intervaltree import IntervalTree
from collections import defaultdict

# Step 1: 寻找重叠 reads（仅限正链）
def find_overlapping_reads(bam_file, min_overlap=100, num_reads=20):
    bam = pysam.AlignmentFile(bam_file, "rb")
    tree = IntervalTree()
    read_positions = {}
    
    # 收集正链 reads 并构建区间树
    for read in bam:
        if not read.is_unmapped and not read.is_supplementary and not read.is_secondary and not read.is_reverse:
            start = read.reference_start
            end = read.reference_end
            if end - start >= min_overlap:  # 确保 read 长度足够
                tree[start:end] = read.query_name
                read_positions[read.query_name] = (start, end)
    
    if len(read_positions) < num_reads:
        raise ValueError(f"Not enough forward strand reads found (found {len(read_positions)}, required {num_reads})")
    
    # 寻找重叠区域并计算覆盖的 reads
    selected_reads = set()
    for interval in tree:
        start, end, query_name = interval.begin, interval.end, interval.data
        overlapping_intervals = tree[start:end]
        if len(overlapping_intervals) >= num_reads:  # 找到包含至少 num_reads 的重叠区域
            overlap_reads = {i.data for i in overlapping_intervals}
            if len(overlap_reads) >= num_reads:
                selected_reads = set(list(overlap_reads)[:num_reads])
                break  # 立即终止
        # 检查重叠区域长度
        for other_interval in overlapping_intervals:
            if other_interval.data != query_name:
                overlap_start = max(start, other_interval.begin)
                overlap_end = min(end, other_interval.end)
                if overlap_end - overlap_start >= min_overlap:
                    selected_reads.add(query_name)
                    selected_reads.add(other_interval.data)
                    if len(selected_reads) >= num_reads:
                        break  # 提前终止内层循环
            if len(selected_reads) >= num_reads:
                break  # 提前终止外层循环
        if len(selected_reads) >= num_reads:
            break  # 提前终止最外层循环
    
    if len(selected_reads) < num_reads:
        print(f"Warning: Only {len(selected_reads)} reads found with sufficient overlap.")
    
    return selected_reads

# Step 2: 从 POD5 文件中提取选定的 reads
def create_new_pod5(input_path, output_pod5, read_ids, is_directory=False):
    # 检查输出文件是否已存在
    if os.path.exists(output_pod5):
        print(f"Output file {output_pod5} already exists. Overwriting it.")
        os.remove(output_pod5)
    
    with Pod5Writer(output_pod5) as writer:
        if is_directory:
            # 处理文件夹中的所有 POD5 文件
            for pod5_file in os.listdir(input_path):
                if pod5_file.endswith(".pod5"):
                    pod5_path = os.path.join(input_path, pod5_file)
                    with pod5.Reader(pod5_path) as reader:
                        for read in reader.reads():
                            if str(read.read_id) in read_ids:
                                r = pod5.Read(
                                    read_id=read.read_id,
                                    read_number=read.read_number,
                                    start_sample=read.start_sample,
                                    median_before=read.median_before,
                                    end_reason=read.end_reason,
                                    calibration=read.calibration,
                                    pore=read.pore,
                                    signal=read.signal,
                                    run_info=read.run_info
                                )
                                writer.add_read(r)
        else:
            # 处理单个 POD5 文件
            with pod5.Reader(input_path) as reader:
                for read in reader.reads():
                    if str(read.read_id) in read_ids:
                        r = pod5.Read(
                            read_id=read.read_id,
                            read_number=read.read_number,
                            start_sample=read.start_sample,
                            median_before=read.median_before,
                            end_reason=read.end_reason,
                            calibration=read.calibration,
                            pore=read.pore,
                            signal=read.signal,
                            run_info=read.run_info
                        )
                        writer.add_read(r)
    print(f"New POD5 file created: {output_pod5}")

# Step 3: 设置 argparse 并运行主程序
def main():
    parser = argparse.ArgumentParser(description="Extract overlapping forward-strand reads from POD5 file(s) and create a new POD5 file.")
    parser.add_argument("-i", "--input", required=True, help="Input POD5 file or directory containing POD5 files.")
    parser.add_argument("-b", "--bam", required=True, help="Input BAM file with read alignments.")
    parser.add_argument("-o", "--output", required=True, help="Output POD5 file path.")
    parser.add_argument("-m", "--min-overlap", type=int, default=100, help="Minimum overlap length (default: 100).")
    parser.add_argument("-n", "--num-reads", type=int, default=20, help="Number of overlapping reads to select (default: 20).")
    
    args = parser.parse_args()
    
    # 检查输入路径是文件还是文件夹
    is_directory = os.path.isdir(args.input)
    if is_directory and not os.path.exists(args.input):
        raise FileNotFoundError(f"Input directory {args.input} does not exist.")
    elif not is_directory and not os.path.isfile(args.input):
        raise FileNotFoundError(f"Input file {args.input} does not exist.")
    
    # 检查 BAM 文件是否存在
    if not os.path.isfile(args.bam):
        raise FileNotFoundError(f"BAM file {args.bam} does not exist.")
    
    # 确保输出目录存在
    output_dir = os.path.dirname(args.output)
    if output_dir and not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    # 筛选重叠 reads（仅正链）
    overlapping_read_ids = find_overlapping_reads(args.bam, args.min_overlap, args.num_reads)
    if not overlapping_read_ids:
        print("No overlapping forward-strand reads found.")
        return
    
    print(f"Selected forward-strand read IDs: {overlapping_read_ids}")
    
    # 创建新的 POD5 文件
    create_new_pod5(args.input, args.output, overlapping_read_ids, is_directory)

if __name__ == "__main__":
    main()